import json
import os
import datasets
import numpy as np

out_path = '/home/ljw22/workspace/qwen/MUSER-main/predicate_72B_top100.json'
# 训练集和数据集合并，因为不需要进行训练，全部进行测试
train_test_path = '/home/ljw22/workspace/qwen/MUSER-main/data/cases/train_test.json'
train_test_file = open(train_test_path, 'r')
train_test = json.load(train_test_file)
test_querys = train_test['train'] + train_test['test']  #


qc_pairs_path = '/home/ljw22/workspace/qwen/MUSER-main/data/cases/cands_by_query.json'
qc_pairs_file = open(qc_pairs_path, 'r')
qc_pairs = json.load(qc_pairs_file)

feature_path = "/home/ljw22/workspace/qwen/MUSER-main/data/features/72B_feature.json"

with open(feature_path) as fin:
    feature = json.load(fin)

docid_list = list(feature.keys())

docid_id_dict = {}

for id, docid in enumerate(docid_list):
    docid_id_dict[docid] = id


feature_matric = []
for doc in docid_list:
    feature_matric.append(feature[doc])  

feature_array = np.array(feature_matric)

print("计算相似度")
feature_score = np.dot(feature_array, feature_array.transpose())
print("相似度排序")
sorted_ndarray = np.argsort(feature_score, axis=1)
sorted_array = sorted_ndarray.tolist()

predicate_top100 = {}

print("确定候选")
for idx, cand in enumerate(sorted_array):
    docid = docid_list[idx]
    if docid not in qc_pairs.keys():
        continue
    
    cand_id_list = []
    for c in cand:
        cid = str(docid_list[c])
        if int(cid) in qc_pairs[docid]:
            cand_id_list.append(cid)
    
    assert len(cand_id_list) == 100
    predicate_top100[docid] = cand_id_list


# print(predicate_top100)
out_file = open(out_path, 'w')
json.dump(predicate_top100, out_file)
